-- Tencent is pleased to support the open source community by making BK-CI 蓝鲸持续集成平台 available.
-- Copyright (C) 2019 THL A29 Limited, a Tencent company.  All rights reserved.
-- BK-CI 蓝鲸持续集成平台 is licensed under the MIT license.
-- A copy of the MIT License is included in this file.
-- Terms of the MIT License:
-- ---------------------------------------------------
-- Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
-- documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
-- rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
-- permit persons to whom the Software is furnished to do so, subject to the following conditions:
-- The above copyright notice and this permission notice shall be included in all copies or substantial portions of
-- the Software.
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
-- LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
-- NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
-- WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
-- SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
local _M = {}
function _M:isAccess()
    -- report服务和bkrepo服务不做频率限制
    if ngx.var.service == 'report' or ngx.var.service == 'bkrepo' or ngx.var.service == 'schedule' or ngx.var.service ==
        'task' then
        return true
    end

    local isAccess = true

    if ngx.var.access_type == 'user' then -- 用户请求类型访问频率限制
        isAccess = self:isUserAccess()
    elseif ngx.var.access_type == 'build' or ngx.var.access_type == 'external' then -- ip地址请求类型访问频率限制
        isAccess = self:isIpAccess()
    end

    return isAccess
end

function _M.isIpAccess()
    local limit_req = require "resty.limit.req"
    -- 创建req_limit实例，每秒100次请求，100个等待，超过200个的请求全部拒绝
    -- 外部请求频率50,浮动20；构建机类请求频率20，浮动20
    local lim = nil
    if ngx.var.access_type == 'external' then
        lim, err = limit_req.new("build_limit_req_store", 100, 50)
    else
        lim, err = limit_req.new("build_limit_req_store", 100, 20)
    end
    -- 创建req_limit实例失败时
    if not lim then
        ngx.log(ngx.STDERR, "failed to instantiate a resty.limit.req object: ", err)
        return false
    end
    -- 获取4字节的IP的KEY
    local key = ngx.var.http_x_devops_real_ip or ngx.var.http_x_real_ip or ngx.var.binary_remote_addr or "0.0.0.0"
    -- 获取key目前的状态:delay非空的时候，说明接受请求，err是排队信息；delay为空的时候，说明拒绝请求，err是错误信息。
    local delay, err = lim:incoming(key, true)
    if not delay then
        if err == "rejected" then
            return false
        end
        ngx.log(ngx.STDERR, "failed to limit req: ", err)
        return false
    end
    if delay >= 0.001 then
        -- 排队号
        local excess = err
        -- 等待需要delay的时间
        ngx.sleep(delay)
        return true
    end
    --- 不用等待
    return true
end

function _M.isUserAccess()
    local limit_req = require "resty.limit.req"
    -- 创建req_limit实例，每秒20请求，20个等待，超过40个的请求全部拒绝
    local lim, err = limit_req.new("user_limit_req_store", 20, 20)
    -- 创建req_limit实例失败时
    if not lim then
        ngx.log(ngx.STDERR, "failed to instantiate a resty.limit.req object: ", err)
        return false
    end
    -- 获取4字节的IP的KEY
    local key = ngx.var.uid
    -- 获取key目前的状态:delay非空的时候，说明接受请求，err是排队信息；delay为空的时候，说明拒绝请求，err是错误信息。
    local delay, err = lim:incoming(key, true)
    if not delay then
        if err == "rejected" then
            return false
        end
        ngx.log(ngx.STDERR, "failed to limit req: ", err)
        return false
    end
    if delay >= 0.001 then
        -- 排队号
        local excess = err
        -- 等待需要delay的时间
        ngx.sleep(delay)
        return true
    end

    --- 不用等待
    return true
end

return _M
